package main

import (
	"MNIST-CNN/pkg/network"
	"MNIST-CNN/pkg/training"
	"fmt"
	"log"
)

func main() {
	// 加载数据集
	trainDataset, testDataset, err := training.LoadDataset()
	if err != nil {
		log.Fatalf("加载数据集失败: %v", err)
	}

	// 打印数据集信息
	fmt.Printf("训练数据集包含 %d 个样本\n", len(trainDataset.Images))
	fmt.Printf("测试数据集包含 %d 个样本\n", len(testDataset.Images))

	// 图像尺寸 (MNIST为28x28)
	imageHeight := 28
	imageWidth := 28
	numClasses := 10

	// 创建神经网络模型
	cnn := network.NewNeuralNetwork()

	// 添加卷积层 - 6个5x5的滤波器
	convLayer1 := network.NewConvLayer(1, imageHeight, imageWidth, 6, 5, 5, 1, 0)
	cnn.AddLayer(convLayer1)

	// 添加池化层 - 2x2的最大池化
	poolLayer1 := network.NewPoolLayer(6, 24, 24, 2, 2, 2, network.Max)
	cnn.AddLayer(poolLayer1)

	// 添加第二个卷积层 - 16个5x5的滤波器
	convLayer2 := network.NewConvLayer(6, 12, 12, 16, 5, 5, 1, 0)
	cnn.AddLayer(convLayer2)

	// 添加第二个池化层 - 2x2的最大池化
	poolLayer2 := network.NewPoolLayer(16, 8, 8, 2, 2, 2, network.Max)
	cnn.AddLayer(poolLayer2)

	// 添加扁平化层，将3D张量转换为1D向量
	flattenLayer := network.NewFlattenLayer(16, 4, 4)
	cnn.AddLayer(flattenLayer)

	// 添加全连接层
	fcLayer1 := network.NewFCLayer(256, 120, network.ReLU, network.ReLUDerivative)
	cnn.AddLayer(fcLayer1)

	// 添加第二个全连接层
	fcLayer2 := network.NewFCLayer(120, 84, network.ReLU, network.ReLUDerivative)
	cnn.AddLayer(fcLayer2)

	// 添加输出层
	outputLayer := network.NewFCLayer(84, numClasses, network.Softmax, nil)
	cnn.AddLayer(outputLayer)

	// 训练参数
	batchSize := 64      // mini-batch大小
	learningRate := 0.01 // 学习率
	epochs := 10         // 训练轮数

	// 训练模型
	training.TrainCNN(cnn, trainDataset, testDataset, batchSize, learningRate, epochs, numClasses)

	// 展示一些测试样本的预测结果
	testInputs := network.PrepareImageBatch(testDataset.Images[:10], imageHeight, imageWidth)
	for i := 0; i < 10; i++ {
		prediction := cnn.Predict(testInputs[i])
		fmt.Printf("样本 %d 的预测类别：%d, 真实类别：%d\n", i+1, prediction, testDataset.Labels[i])
	}
}
